from flask import Flask, render_template, request, jsonify, send_from_directory
import os
import time
import subprocess
import re
from functools import wraps
import logging
from logging.handlers import RotatingFileHandler
from datetime import datetime
import threading
from collections import defaultdict
import json

app = Flask(__name__)

# 配置日志
if not os.path.exists('logs'):
    os.makedirs('logs')
logging.basicConfig(
    handlers=[RotatingFileHandler('logs/app.log', maxBytes=100000, backupCount=3)],
    level=logging.INFO,
    format='%(asctime)s %(levelname)s: %(message)s [in %(pathname)s:%(lineno)d]'
)

# 配置音频文件存储目录
AUDIO_FOLDER = 'audio_files'
if not os.path.exists(AUDIO_FOLDER):
    os.makedirs(AUDIO_FOLDER)

# 安全配置
MAX_CONTENT_LENGTH = 1024 * 1024  # 1MB
ALLOWED_EXTENSIONS = {'wav'}
DOCKER_CONTAINER = os.getenv('DOCKER_CONTAINER', 'voice2text')
DOCKER_PATH = os.getenv('DOCKER_PATH', '/root/uploadfiles')
REQUEST_LIMIT = 30  # 每分钟最大请求次数
request_counts = defaultdict(list)
request_lock = threading.Lock()

# 格式有：318641893286795_atu43738609, 1736926681_at-ts, ios_usrid_1736931943_P5_0, asid_1736927370551_P1_2
def is_valid_id(file_id):
    """验证 ID 格式是否合法"""
    return bool(re.match(r'^[a-zA-Z0-9_-]+$', file_id))

def rate_limit(f):
    """请求频率限制装饰器"""
    @wraps(f)
    def decorated_function(*args, **kwargs):
        current_time = time.time()
        client_ip = request.remote_addr

        with request_lock:
            # 清理过期的请求记录（1分钟前的）
            request_counts[client_ip] = [t for t in request_counts[client_ip]
                                       if current_time - t < 60]

            # 检查请求频率
            if len(request_counts[client_ip]) >= REQUEST_LIMIT:
                app.logger.warning(f'Rate limit exceeded for IP: {client_ip}')
                return jsonify({'status': 'error',
                              'message': '请求过于频繁，请稍后再试'}), 429

            request_counts[client_ip].append(current_time)

        return f(*args, **kwargs)
    return decorated_function

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/search', methods=['POST'])
@rate_limit
def search():
    try:
        if not request.is_json:
            return jsonify({'status': 'error', 'message': '无效的请求格式'})

        data = request.get_json()
        id_query = data.get('id', '').strip()
        text_query = data.get('query', '').strip()

        if not id_query and not text_query:
            return jsonify({'status': 'error', 'message': '搜索内容不能为空'})

        if id_query and not is_valid_id(id_query):
            app.logger.warning(f'Invalid id: {id_query}')
            return jsonify({'status': 'error', 'message': '无效的 ID 格式'})

        # 清理 audio_files 目录
        for file in os.listdir(AUDIO_FOLDER):
            file_path = os.path.join(AUDIO_FOLDER, file)
            if os.path.isfile(file_path):
                os.remove(file_path)

        if id_query:
            # 可能有多个
            # 先在 Docker 容器中查找最近一天修改的 wav 文件
            # 因为 sid 可能会重复，所以每个 sid 后面加了时间戳，wav 文件名就是: sid_时间戳.wav
            # 又因为 docker 命令不支持通配符, 所以只能先把文件查询出来, 然后再找出包含 sid 的文件
            find_cmd = ['docker', 'exec', DOCKER_CONTAINER, 'find', DOCKER_PATH, '-ctime', '-1']
            try:
                result = subprocess.run(find_cmd, capture_output=True, text=True, timeout=10)
                if result.returncode != 0:
                    app.logger.error(f"Docker find error: {result.stderr}")
                    return jsonify({'status': 'error', 'message': '查找音频文件失败'})

                # 从结果中筛选包含 ID 的文件
                wav_files = []
                for line in result.stdout.splitlines():
                    if line.endswith('.wav') and id_query in line:
                        wav_files.append(os.path.basename(line))

                if not wav_files:
                    return jsonify({'status': 'not_found', 'message': f'未找到该 ID 的音频数据'})

                # 复制找到的文件到本地
                valid_files = []
                for wav_file in wav_files:
                    cp_cmd = ['docker', 'cp', f'{DOCKER_CONTAINER}:{DOCKER_PATH}/{wav_file}', os.path.join(AUDIO_FOLDER, wav_file)]
                    result = subprocess.run(cp_cmd, capture_output=True, text=True, timeout=10)

                    if result.returncode != 0:
                        app.logger.error(f"Docker cp error for {wav_file}: {result.stderr}")
                        continue

                    wav_path = os.path.join(AUDIO_FOLDER, wav_file)
                    if os.path.exists(wav_path):
                        if os.path.getsize(wav_path) > MAX_CONTENT_LENGTH:
                            os.remove(wav_path)
                            app.logger.warning(f'文件 {wav_file} 超出大小限制，已删除')
                        else:
                            valid_files.append(wav_file)

                if not valid_files:
                    return jsonify({'status': 'error', 'message': '所有文件都超出大小限制, 只支持 1MB 以内的文件'})

                return jsonify({'status': 'success', 'files': valid_files})

            except subprocess.TimeoutExpired:
                app.logger.error("Docker command timeout")
                return jsonify({'status': 'error', 'message': '操作超时'})
        else:
            # 文本搜索：只支持获取最新的一个
            # 验证文本内容
            if len(text_query) > 100:  # 限制查询文本长度
                return jsonify({'status': 'error', 'message': '查询文本过长'})

            app.logger.info(f"使用文本搜索: {text_query}")
            result = subprocess.run(['python', './process_query.py', text_query],
                                capture_output=True,
                                text=True,
                                timeout=30)  # 添加超时限制

            if result.returncode != 0:
                app.logger.error(f"Process query error: {result.stderr}")
                return jsonify({'status': 'error', 'message': '处理失败'})

            files = [f for f in os.listdir(AUDIO_FOLDER)
                    if os.path.isfile(os.path.join(AUDIO_FOLDER, f))]

            if not files:
                return jsonify({'status': 'not_found', 'message': '未找到相关音频数据'})

            # 检查所有文件的大小并收集合法的文件
            valid_files = []
            for wav_file in files:
                wav_path = os.path.join(AUDIO_FOLDER, wav_file)
                if os.path.getsize(wav_path) > MAX_CONTENT_LENGTH:
                    os.remove(wav_path)
                    app.logger.warning(f'文件 {wav_file} 超出大小限制，已删除')
                else:
                    valid_files.append(wav_file)

            if not valid_files:
                return jsonify({'status': 'error', 'message': '所有文件都超出大小限制, 只支持 1MB 以内的文件'})

            return jsonify({'status': 'success', 'files': valid_files})
    except subprocess.TimeoutExpired:
        app.logger.error("Operation timeout")
        return jsonify({'status': 'error', 'message': '操作超时'})
    except Exception as e:
        app.logger.error(f"Error: {str(e)}")
        return jsonify({'status': 'error', 'message': f'服务器内部错误: {str(e)} 应该是没有 docker'})

@app.route('/audio/<filename>')
@rate_limit
def serve_audio(filename):
    # 验证文件名
    if not is_valid_id(os.path.splitext(filename)[0]):
        app.logger.warning(f'Invalid filename attempted: {filename}')
        return jsonify({'status': 'error', 'message': '无效的文件名'}), 400

    # 验证文件扩展名
    if not filename.lower().endswith('.wav'):
        return jsonify({'status': 'error', 'message': '无效的文件类型'}), 400

    return send_from_directory(AUDIO_FOLDER, filename)

if __name__ == '__main__':
    # 需要的话，可以将 debug 设置为 True
    app.run(debug=False, host='0.0.0.0', port=9111)
